"""
author：石沙
date：2020-09-28
content：本模块用进行文本增强和下采样
"""

# 如下导入时为保证训练时的任务流能正常执行
import sys
from settings import MAIN_PATH, SRC_PATH
sys.path.extend([MAIN_PATH, SRC_PATH])


from site_packages.ml_libs.nlpkits import eda
import pandas as pd
from datasets import load_book_clean, load_book_augmented
from site_packages.utils.models import ModelOp


def EDA():
    # 设定要进行数据增强的标签
    label_to_augment = [0, 1, 2, 3]

    # 加载数据
    data = load_book_clean()
    data_to_augment = data[['book_id', 'label', 'full_content']][data['label'].isin(label_to_augment)]
    data_to_augment['book_id'] = data_to_augment['book_id'].astype(str)

    # # 训练时候用来保存和加载中间结果
    # ModelOp.save(data_to_augment, 'data_to_augment', is_model=False)
    # data_to_augment = ModelOp.load_data('data_to_augment')

    # 设定要用的EDA方法
    transformers = [
        eda.SynonymsReplacement(),
        eda.RandomInsert(),
        eda.RandomSwap(),
        eda.RandomDeletion()
    ]

    # 执行EDA
    new_texts = eda.augment_by_eda(
        data_to_augment['full_content'],
        data_to_augment['book_id'],
        data_to_augment['label'],
        id_column='book_id',
        label_column='label',
        transformers=transformers
    )

    # 保存
    new_texts.rename(columns={'eda_text': 'full_content'}, inplace=True)
    data_augmented = pd.concat([data_to_augment, new_texts])
    data_augmented = data_augmented.sample(frac=1, random_state=1)
    ModelOp.save(data_augmented, 'book_augmented', is_model=False)


def random_under_sample(data, label_column='label', num_per_label=None):
    label_counts = data[label_column].value_counts()
    min_counts = num_per_label if num_per_label is not None else label_counts.min()
    data_sampled = []
    for label in label_counts.index:
        data_by_label = data[data[label_column] == label]
        sampling_percent = min_counts / data_by_label.shape[0]
        data_sampled.append(data_by_label.sample(frac=sampling_percent))
    return pd.concat(data_sampled).sample(frac=1)


def under_sampling(target='normal', file_name=None, num_per_label=None):
    """
    :param target: normal，针对预处理后的原数据集；augment，针对EDA增强后的数据集
    :param file_name: 指定下采样后保存的文件名称
    :param num_per_label: 指定每类标签的采样数量，如果为None，则采用记录数最小的label
        对应的记录数
    """
    assert file_name is not None
    if target == 'normal':
        data = load_book_clean()
    else:
        data = load_book_augmented()
    data = data[['book_id', 'label', 'full_content']].copy()
    data = random_under_sample(data, num_per_label=num_per_label)
    ModelOp.save(data, file_name, is_model=False)


if __name__ == '__main__':
    EDA()
    under_sampling(target='normal', file_name='undersample_1k', num_per_label=1000)


